__author__ = 'pcJxy'

from utils.prime_utils import gen_prime, is_prime
from utils.codec_utils import int_2_n_bytes, b64str_2_int, bytes_2_int
import sys


def quick_pow(base, exp, mod):
    t = 1
    while exp != 0:
        if exp & 1 == 1:
            t = (t % mod * base % mod) % mod
        exp >>= 1
        base = (base * base) % mod
    return t


def extended_euclidean(a, b):
    if b == 0:
        return 1, 0
    xx, yy = extended_euclidean(b, a % b)
    x = yy
    y = xx - (a // b) * yy
    return x, y


def rsa_keygen(p=-1, q=-1, bit_width=1024) -> (int, int, int):
    """
    Generate RSA keys.

    :param p:
    :param q:
    :param bit_width:
    :return: a tuple (e, n, d) which are public key, modulo number and private key respectively.
    """
    p = p if p != -1 else gen_prime(bit_width)
    q = q if q != -1 else gen_prime(bit_width)
    if not is_prime(p) or not is_prime(q):
        print("p or q is not a prime number, please check.", file=sys.stderr)
        return None
    n = p * q
    phi_n = (p - 1) * (q - 1)
    e = 17
    d, _ = extended_euclidean(e, phi_n)
    d = (d + phi_n) % phi_n  # We make sure that d is a positive number.
    return e, n, d


def __rsa_read_file(filename):
    """
    Read a file of base64 encoded strings.

    :param filename:
    :return: decoded integer of the base64 string.
    """
    with open(filename, "r") as infile:
        b64str = "".join(infile.readlines())
        return b64str_2_int(b64str)


def rsa_read_public_key(public_key_filename='./rsa_public.key', n_filename='./rsa_n.key'):
    return __rsa_read_file(public_key_filename), __rsa_read_file(n_filename)


def rsa_read_private_key(private_key_filename='./rsa_private.key', n_filename='./rsa_n.key'):
    return __rsa_read_file(private_key_filename), __rsa_read_file(n_filename)


def __rsa_n_bit_length(n: int):
    """
    Find the first bit length not less than the bit length of `n`.

    :param n:
    :return: a minimum bit length larger than the bit length of `n`
    """
    n_bit_length = 1
    while n_bit_length < n.bit_length():
        n_bit_length <<= 1
    return n_bit_length


def rsa_encrypt(s: bytes, key: int, n: int):
    n_bit_length = __rsa_n_bit_length(n)
    encoded_sequence = []
    for byte in s:
        num = int(byte)
        powered_num = quick_pow(num, key, n)
        encoded_sequence.append(int_2_n_bytes(powered_num, n_bit_length))
    return b''.join(encoded_sequence)


def rsa_decrypt(s: bytes, key: int, n: int):
    n_bit_length = __rsa_n_bit_length(n)
    decoded_sequence = []
    for i in range(0, len(s), n_bit_length):
        encrypted_subsequence = s[i:i + n_bit_length]
        decrypted_byte = quick_pow(bytes_2_int(encrypted_subsequence), key, n)
        decoded_sequence.append(int_2_n_bytes(decrypted_byte, 1))
    return b''.join(decoded_sequence)
